/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 */
package org.openhitls.tls;

import org.junit.BeforeClass;
import org.junit.Test;
import static org.junit.Assert.*;

import javax.net.ssl.*;
import java.io.File;
import java.nio.ByteBuffer;
import java.security.*;
import java.security.cert.X509Certificate;

import org.openhitls.crypto.jce.provider.HiTls4jProvider;

/**
 * Unit tests for TLS/SSL functionality using HiTLS JSSE provider.
 */
public class SSLContextTest {

    private static final String PROVIDER_NAME = HiTls4jProvider.PROVIDER_NAME;

    @BeforeClass
    public static void loadNativeLibraries() {
        String libraryPath = new File(System.getProperty("user.dir") + "/target/native/libhitls_crypto_jni.so").getAbsolutePath();
        System.load(libraryPath);
        Security.addProvider(new HiTls4jProvider());
    }

    // ==================== SSLContext Tests ====================

    @Test
    public void testSSLContextTLSAvailable() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        assertNotNull("SSLContext.TLS should be available", ctx);
        assertEquals("Provider should be HiTLS", PROVIDER_NAME, ctx.getProvider().getName());
    }

    @Test
    public void testSSLContextTLS12Available() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLSv1.2", PROVIDER_NAME);
        assertNotNull("SSLContext.TLSv1.2 should be available", ctx);
    }

    @Test
    public void testSSLContextTLS13Available() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLSv1.3", PROVIDER_NAME);
        assertNotNull("SSLContext.TLSv1.3 should be available", ctx);
    }

    @Test
    public void testSSLContextInitWithNullManagers() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        // Initialize with null managers (should use defaults)
        ctx.init(null, null, null);
        // Should not throw exception
    }

    @Test
    public void testSSLContextInitWithTrustAllManager() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);

        // Create a trust-all TrustManager
        TrustManager[] trustAllCerts = new TrustManager[] {
            new X509TrustManager() {
                public X509Certificate[] getAcceptedIssuers() { return new X509Certificate[0]; }
                public void checkClientTrusted(X509Certificate[] certs, String authType) {}
                public void checkServerTrusted(X509Certificate[] certs, String authType) {}
            }
        };

        ctx.init(null, trustAllCerts, new SecureRandom());
        assertNotNull("SSLContext should be initialized", ctx);
    }

    // ==================== SSLEngine Tests ====================

    @Test
    public void testCreateSSLEngine() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();
        assertNotNull("SSLEngine should be created", engine);
    }

    @Test
    public void testCreateSSLEngineWithHostPort() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine("localhost", 8443);
        assertNotNull("SSLEngine should be created", engine);
        assertEquals("Peer host should be set", "localhost", engine.getPeerHost());
        assertEquals("Peer port should be set", 8443, engine.getPeerPort());
    }

    @Test
    public void testSSLEngineClientMode() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);

        // Default should allow setting client mode
        engine.setUseClientMode(true);
        assertTrue("Should be in client mode", engine.getUseClientMode());

        // Create another engine for server mode
        SSLEngine serverEngine = ctx.createSSLEngine();
        serverEngine.setUseClientMode(false);
        assertFalse("Should be in server mode", serverEngine.getUseClientMode());
    }

    @Test
    public void testSSLEngineHandshakeStatus() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);

        // Before handshake begins
        SSLEngineResult.HandshakeStatus status = engine.getHandshakeStatus();
        assertEquals("Initial status should be NOT_HANDSHAKING",
                     SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING, status);
    }

    @Test
    public void testSSLEngineBeginHandshake() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        // After beginHandshake, client should need to wrap
        SSLEngineResult.HandshakeStatus status = engine.getHandshakeStatus();
        assertEquals("Client should need to wrap after beginHandshake",
                     SSLEngineResult.HandshakeStatus.NEED_WRAP, status);
    }

    @Test
    public void testSSLEngineSession() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);

        SSLSession session = engine.getSession();
        assertNotNull("Session should not be null", session);

        // Session should have reasonable buffer sizes
        assertTrue("App buffer size should be positive",
                   session.getApplicationBufferSize() > 0);
        assertTrue("Packet buffer size should be positive",
                   session.getPacketBufferSize() > 0);
    }

    @Test
    public void testSSLEngineCipherSuites() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();

        // Get supported cipher suites
        String[] supported = engine.getSupportedCipherSuites();
        assertNotNull("Supported cipher suites should not be null", supported);
        assertTrue("Should have at least one supported cipher suite", supported.length > 0);

        // Get enabled cipher suites
        String[] enabled = engine.getEnabledCipherSuites();
        assertNotNull("Enabled cipher suites should not be null", enabled);
        assertTrue("Should have at least one enabled cipher suite", enabled.length > 0);
    }

    @Test
    public void testSSLEngineProtocols() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();

        // Get supported protocols
        String[] supported = engine.getSupportedProtocols();
        assertNotNull("Supported protocols should not be null", supported);
        assertTrue("Should support at least one protocol", supported.length > 0);

        // Get enabled protocols
        String[] enabled = engine.getEnabledProtocols();
        assertNotNull("Enabled protocols should not be null", enabled);
        assertTrue("Should have at least one enabled protocol", enabled.length > 0);

        // Check for TLS 1.2 or 1.3 support
        boolean hasTLS12or13 = false;
        for (String protocol : supported) {
            if (protocol.equals("TLSv1.2") || protocol.equals("TLSv1.3")) {
                hasTLS12or13 = true;
                break;
            }
        }
        assertTrue("Should support TLSv1.2 or TLSv1.3", hasTLS12or13);
    }

    @Test
    public void testSSLEngineClientAuth() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();
        engine.setUseClientMode(false); // Server mode

        // Test needClientAuth
        assertFalse("needClientAuth should be false by default", engine.getNeedClientAuth());
        engine.setNeedClientAuth(true);
        assertTrue("needClientAuth should be true after setting", engine.getNeedClientAuth());

        // Test wantClientAuth
        engine.setNeedClientAuth(false);
        assertFalse("wantClientAuth should be false by default", engine.getWantClientAuth());
        engine.setWantClientAuth(true);
        assertTrue("wantClientAuth should be true after setting", engine.getWantClientAuth());
    }

    @Test
    public void testSSLEngineCloseOperations() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();
        engine.setUseClientMode(true);

        // Initially not done
        assertFalse("Inbound should not be done initially", engine.isInboundDone());
        assertFalse("Outbound should not be done initially", engine.isOutboundDone());

        // Close outbound
        engine.closeOutbound();
        assertTrue("Outbound should be done after closeOutbound", engine.isOutboundDone());

        // Close inbound
        engine.closeInbound();
        assertTrue("Inbound should be done after closeInbound", engine.isInboundDone());
    }

    @Test
    public void testSSLEngineSetEnabledCipherSuites() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();

        // Get supported and set a subset as enabled
        String[] supported = engine.getSupportedCipherSuites();
        if (supported.length > 1) {
            String[] subset = new String[] { supported[0] };
            engine.setEnabledCipherSuites(subset);

            String[] enabled = engine.getEnabledCipherSuites();
            assertEquals("Should have only one enabled cipher suite", 1, enabled.length);
            assertEquals("Enabled suite should match", supported[0], enabled[0]);
        }
    }

    @Test
    public void testSSLEngineSetEnabledProtocols() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();

        // Set only TLSv1.2
        engine.setEnabledProtocols(new String[] { "TLSv1.2" });
        String[] enabled = engine.getEnabledProtocols();
        assertEquals("Should have only one enabled protocol", 1, enabled.length);
        assertEquals("Enabled protocol should be TLSv1.2", "TLSv1.2", enabled[0]);
    }

    @Test
    public void testSSLEngineDelegatedTask() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();
        engine.setUseClientMode(true);

        // HiTLS doesn't require delegated tasks
        Runnable task = engine.getDelegatedTask();
        assertNull("Delegated task should be null for HiTLS", task);
    }

    @Test
    public void testSSLEngineSSLParameters() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine("example.com", 443);
        engine.setUseClientMode(true);

        SSLParameters params = engine.getSSLParameters();
        assertNotNull("SSLParameters should not be null", params);

        // Verify parameters are populated
        assertNotNull("Cipher suites should be set", params.getCipherSuites());
        assertNotNull("Protocols should be set", params.getProtocols());
    }

    // ==================== Factory Tests ====================

    @Test
    public void testSSLSocketFactory() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLSocketFactory factory = ctx.getSocketFactory();
        assertNotNull("SSLSocketFactory should not be null", factory);

        // Verify default cipher suites
        String[] defaultCiphers = factory.getDefaultCipherSuites();
        assertNotNull("Default cipher suites should not be null", defaultCiphers);
        assertTrue("Should have default cipher suites", defaultCiphers.length > 0);
    }

    @Test
    public void testSSLServerSocketFactory() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLServerSocketFactory factory = ctx.getServerSocketFactory();
        assertNotNull("SSLServerSocketFactory should not be null", factory);

        // Verify supported cipher suites
        String[] supported = factory.getSupportedCipherSuites();
        assertNotNull("Supported cipher suites should not be null", supported);
        assertTrue("Should have supported cipher suites", supported.length > 0);
    }

    // ==================== Session Context Tests ====================

    @Test
    public void testClientSessionContext() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLSessionContext sessionCtx = ctx.getClientSessionContext();
        assertNotNull("Client session context should not be null", sessionCtx);
    }

    @Test
    public void testServerSessionContext() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLSessionContext sessionCtx = ctx.getServerSessionContext();
        assertNotNull("Server session context should not be null", sessionCtx);
    }

    // ==================== Negative Tests ====================

    @Test(expected = IllegalStateException.class)
    public void testSSLContextNotInitializedCreateEngine() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        // Do not initialize
        ctx.createSSLEngine(); // Should throw IllegalStateException
    }

    @Test(expected = IllegalStateException.class)
    public void testSSLContextNotInitializedGetSocketFactory() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        // Do not initialize
        ctx.getSocketFactory(); // Should throw IllegalStateException
    }

    @Test(expected = IllegalArgumentException.class)
    public void testSSLEngineSetClientModeAfterHandshake() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();
        engine.setUseClientMode(true);
        engine.beginHandshake();

        // Try to change mode after handshake started - should throw
        engine.setUseClientMode(false);
    }

    @Test
    public void testSSLEngineMultipleBeginHandshake() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);
        ctx.init(null, null, null);

        SSLEngine engine = ctx.createSSLEngine();
        engine.setUseClientMode(true);
        engine.beginHandshake();

        // Calling beginHandshake again during handshake is implementation-defined
        // Our implementation allows it (transitions state appropriately)
        try {
            engine.beginHandshake();
            // If it doesn't throw, that's acceptable for our implementation
        } catch (SSLException e) {
            // If it throws SSLException, that's also acceptable (renegotiation denied)
        }
        // Either behavior is acceptable
    }
}
